import json

import numpy as np
import pymysql
import torch
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder, MinMaxScaler

from EnergyConfig import EnergyConfig


class StockDataProvider:

    def __init__(self):
        self.mysql_host = "127.0.0.1"
        self.mysql_port = 3306

        self.connect = pymysql.connect(
            host=self.mysql_host,
            port=self.mysql_port,
            user='root',
            passwd='hxsoft.net',
            database='stock')
        self.cursor = self.connect.cursor()

    def get_stock_data(self):
        """
        获取原始数据
        :return:    原始数据Tensor
        """
        try:
            self.cursor.execute("select * from stock_result where stock_high != ''")
            rows = self.cursor.fetchall()
            if rows is not None:
                seq_list = []
                value_list = []
                for row in rows:
                    json_list = json.loads(row[6])
                    item_list = []
                    for index in range(len(json_list)):
                        item = json_list[index]
                        item_list.append(item['open'])
                        item_list.append(item['high'])
                        item_list.append(item['low'])
                        item_list.append(item['close'])
                        item_list.append(item['zl1'])
                        item_list.append(item['zl2'])
                    seq_list.append(item_list)
                    increment = (100 * (row[3] - row[2])) / row[2]
                    if 0 < increment < 3:
                        value_list.append(1)
                    elif 3 <= increment < 6:
                        value_list.append(3)
                    elif 6 <= increment < 10:
                        value_list.append(5)
                    elif increment >= 10:
                        value_list.append(7)
                    elif increment == 0:
                        value_list.append(0)
                    elif -3 < increment < 0:
                        value_list.append(2)
                    elif -6 < increment <= -3:
                        value_list.append(4)
                    elif -10 < increment <= -6:
                        value_list.append(6)
                    else:
                        value_list.append(8)
                sep_data = torch.tensor(seq_list, dtype=torch.float32)
                value_data = torch.tensor(value_list, dtype=torch.long)
                return sep_data, value_data
            else:
                return None
        finally:
            if self.cursor is not None:
                self.cursor.close()
            if self.connect is not None:
                self.connect.close()

    def get_split_data(self):
        """
        DataPart Split
        :return:
        """
        sep_data, value_data = self.get_stock_data()
        if sep_data is not None:
            # for feature_dim in range(StockConfig.n_features):
            scaler = StandardScaler()
            sep_data = scaler.fit_transform(sep_data).reshape(-1, 5, EnergyConfig.n_features).astype(np.float32)

            scaler1 = MinMaxScaler(feature_range=(0, 1))
            value_data = scaler1.fit_transform(value_data.reshape(-1, 1)).reshape(-1, 1).astype(np.float32)

            train_data, test_data, train_value, test_value \
                = train_test_split(sep_data, value_data, test_size=0.2, random_state=42)

            return train_data, test_data, train_value, test_value
        else:
            return None
